/**
 * @file greedy_graph_coloring.cpp
 * @brief
 * 地图染色问题 - 贪心算法
 *
 * @author Wang Bo (1276106221@qq.com)
 * @version 1.0
 * @date 2021-07-09
 *
 * @copyright Wang Bo
 *
 */

#include <iostream>
#include <vector>
#include <string>
#include <map>
#include <unordered_map>
#include <set>
#include <algorithm>
#include <random>
#include <numeric>
#include <cassert>
#include <functional>

using std::cout;
using std::endl;
using std::vector;


template <typename T>
struct Edge {
    size_t src;
    size_t dest;
    T weight;
    inline bool operator<(const Edge<T>& that) const {
        return weight < that.weight;
    }
    inline bool operator>(const Edge<T>& that) const {
        return weight < that.weight;
    }
};

template <typename T>
class Graph {
public:
    Graph(size_t N): vertices_(N) {}
    auto vertices() const { return vertices_; }
    const auto& edges() const { return edgeList_; }
    void addEdge(const Edge<T>& e) {
        if (e.src >= 1 && e.src <= vertices_ && e.dest >= 1 && e.dest <= vertices_)
            edgeList_.emplace_back(e);
        else
            std::cerr << "Edge out of range\n";
    }
    auto edgesFromVertex(size_t v) const {
        vector<Edge<T>> edgesRet;
        for (const auto& e : edgeList_) {
            if (e.src == v)
                edgesRet.emplace_back(e);
        }
        return edgesRet;
    }
    friend std::ostream& operator<<(std::ostream& os, const Graph<T>& G) {
        for (size_t i = 1; i <= G.vertices(); ++i) {
            auto edges = G.edgesFromVertex(i);
            os << "Vertex[" << i << "]:\t";
            for (const auto& e : edges) {
                os << "{" << e.dest << ", " << e.weight << "}, ";
            }
            os << endl;
        }
        return os;
    }
private:
    size_t vertices_; //顶点数
    vector<Edge<T>> edgeList_; //所有边的数组
};

std::unordered_map<size_t, std::string> colorMap {
    {1, "\033[31;1mRed\033[0m"},
    {2, "\033[34;1mBlue\033[0m"},
    {3, "\033[32;1mGreen\033[0m"},
    {4, "\033[33;1mYellow\033[0m"},
    {5, "\033[35;1mPurple\033[0m"},
    {6, "\033[37;1mWhite\033[0m"},
};

template <typename T>
auto greedyColoring(const Graph<T>& G) {
    auto size = G.vertices();
    vector<size_t> assignedColors(size + 1); //已分配的颜色
    for (auto i = 1; i <= size; ++i) { //遍历所有顶点
        auto edgesFromVi = G.edgesFromVertex(i);
        std::set<size_t> neighbourColors;
        for (const auto& e : edgesFromVi) {
            auto destColor = static_cast<size_t>(assignedColors[e.dest]);
            neighbourColors.insert(destColor);
        }

        auto smallestUnassignedColor = 1; //可选颜色初始为1
        for (; smallestUnassignedColor <= colorMap.size(); ++smallestUnassignedColor)
        {
            //遍历所有颜色从中找到邻居中不存在的颜色
            if (neighbourColors.find(smallestUnassignedColor) == neighbourColors.end())
                break;
        }
        //将这个颜色设置为已选
        assignedColors[i] = smallestUnassignedColor;
    }
    return assignedColors;
}

void printColors(const vector<size_t>& colors) {
    for (size_t i = 1; i < colors.size(); ++i) {
        cout << "Vertex: " << (i + 1) << " : " << colorMap[colors[i]] << endl;
    }
}

int main(int argc, const char** argv) {
    using T = size_t;
    Graph<T> G(8);

	std::map<size_t, std::vector<std::pair<size_t, T>>> edges;
	edges[1] = { {2, 2}, {5, 3} };
	edges[2] = { {1, 2}, {5, 5}, {4, 1} };
	edges[3] = { {4, 2}, {7, 3} };
	edges[4] = { {2, 1}, {3, 2}, {5, 2}, {6, 4}, {8, 5} };
	edges[5] = { {1, 3}, {2, 5}, {4, 2}, {8, 3} };
	edges[6] = { {4, 4}, {7, 4}, {8, 1} };
	edges[7] = { {3, 3}, {6, 4} };
	edges[8] = { {4, 5}, {5, 3}, {6, 1} };

	for (auto& i : edges)
		for (auto& j : i.second)
			G.addEdge(Edge<T>{ i.first, j.first, j.second });

	std::cout << "Original Graph: " << std::endl;
	std::cout << G << std::endl;

	auto colors = greedyColoring<T>(G);
	std::cout << "Vertex Colors: " << std::endl;
	printColors(colors);
    return 0;
}